import itertools
import sys

import tables

tmpargv = sys.argv[:]
import ROOT
sys.argv = tmpargv


def _groups_of_n(n, iterable, padvalue=None):
    "groups(3, 'abcdefg', 'x') --> ('a','b','c'), ('d','e','f'), ('g','x','x')"
    return itertools.izip(*[itertools.chain(iterable, itertools.repeat(padvalue, n-1))]*n)


pttypes_to_rootchr = {tables.description.Float32Col: 'F',
                      tables.description.Float64Col: 'D',
                      tables.description.Int32Col: 'I'}

pttypes_to_roottype = {tables.description.Float32Col: 'Float_t',
                       tables.description.Float64Col: 'Double_t',
                       tables.description.Int32Col: 'Int_t'}


class ROOTStruct(object):
    def __init__(self, ptdict, struct_name='MyStruct'):
        self.fields = sorted(ptdict.keys()); 
        self.pttypes = dict([(f, type(ptdict[f])) for f in self.fields])

        self.nper_struct = 10

        self.struct_names = []
        self.which_struct = {}

        for i, tmpfields in enumerate(_groups_of_n(self.nper_struct, self.fields)):
            nm = 'MS{0}'.format(i)
            tmp_ptdict = dict([(f, self.pttypes[f]) for f in tmpfields if f])
            self.struct_names.append(nm)

            for f in tmpfields:
                self.which_struct[f] = nm

            structstr = structstr_from_ptdict(tmp_ptdict, struct_name=nm)
            ROOT.gROOT.ProcessLine(structstr)

        tmpcons = __import__('ROOT', globals(), locals(), self.struct_names, -1)
        
        self.struct_constructors = dict([(nm, getattr(tmpcons, nm)) for nm in self.struct_names])

        self.structs = dict([(nm, self.struct_constructors[nm]()) for nm in self.struct_names])

    def address_of(self, field):
        nm = self.which_struct[field]
        struct = self.structs[nm]
        return ROOT.AddressOf(struct, field)

    def root_typechr(self, field):
        return pttypes_to_rootchr[self.pttypes[field]]
    
    
    def __setitem__(self, field, rhs):
        nm = self.which_struct[field]
        struct = self.structs[nm]
        setattr(struct, field, rhs)
    
    def __getitem__(self, field):
        nm = self.which_struct[field]
        struct = self.structs[nm]
        return getattr(struct, field)
        

        

def structstr_from_ptdict(ptdict, struct_name='MyStruct'):
    fields = sorted(ptdict.keys())

    retstr = "struct {0} {{".format(struct_name)
    for f in fields:    
        ft = ptdict[f]
        retstr += "  {0} {1}; ".format(pttypes_to_roottype[ft], f)
        
    retstr += '};'
                  
    return retstr

        
class BadCommandLineError(Exception):
    def __init__(self, msg, retcode):
        self.msg = msg
        self.retcode = retcode

    def __str__(self):
        return self.msg

